from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain_community.tools.tavily_search import TavilySearchResults

_ = load_dotenv(find_dotenv())
 
model = ChatZhipuAI(
        model="glm-4-plus",
        temperature=0.9,              
    )
#model = ChatOpenAI(model="gpt-4o")

search = TavilySearchResults()
tools = [search]

from langchain_core.prompts import PromptTemplate

template = '''Answer the following questions as best you can. You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}'''

prompt = PromptTemplate.from_template(template)

from langchain.agents import AgentExecutor, create_react_agent
agent = create_react_agent(model, tools, prompt)
agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
)

agent_executor.invoke({"input": "什么是新型电力系统?"})
# agent_executor.invoke({"input": "一年有多少个月份？"})